#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <vector>
#include <map>
#include <string>
#include <list>
#include "block.h"
#include "qcow2.h"
#include "utils.h"

// 全局块设备管理器
static std::map<std::string, BlockBackend*> blk_map;

// 打开文件并注册到管理器
static BlockBackend* open_disk(const char* name, const char* filename) {
    BlockDriverState* bs = bdrv_open(filename, "qcow2");
    if (!bs) {
        report_error("Failed to open file: %s", filename);
        return NULL;
    }

    BlockBackend* blk = blk_new(name);
    if (!blk) {
        bdrv_close(bs);
        return NULL;
    }

    if (blk_attach(blk, bs) != 0) {
        blk_unref(blk);
        bdrv_close(bs);
        return NULL;
    }

    blk_map[std::string(name)] = blk;
    report_info("Opened disk '%s' -> '%s'", name, filename);
    return blk;
}

// 关闭文件并从管理器移除
static void close_disk(const char* name) {
    auto it = blk_map.find(std::string(name));
    if (it != blk_map.end()) {
        blk_unref(it->second);
        blk_map.erase(it);
        report_info("Closed disk: %s", name);
    } else {
        report_error("Disk not found: %s", name);
    }
}

// 获取块后端
static BlockBackend* get_disk(const char* name) {
    auto it = blk_map.find(std::string(name));
    if (it != blk_map.end()) {
        return it->second;
    }
    return NULL;
}

// 列出所有打开的磁盘
static void list_disks() {
    printf("\n=== Managed Disks ===\n");

    if (blk_map.empty()) {
        printf("No disks managed\n");
        return;
    }

    for (const auto& pair : blk_map) {
        BlockBackend* blk = pair.second;
        BlockDriverState* bs = blk_bs(blk);
        if (bs) {
            printf("Disk: %s\n", pair.first.c_str());
            printf("  File: %s\n", bs->filename);
            printf("  Size: %ld MB\n", (bdrv_getlength(bs) / (1024 * 1024)));

            char *backing = qcow2_get_backing_file(bs);
            if (backing) {
                printf("  Backing File: %s\n", backing);
                qemu_free(backing);
            }
        }
    }

    // 也显示所有打开的块设备
    bdrv_list();
}

// 分析文件
static void analyze_file(const char *disk_name) {
    BlockBackend* blk = get_disk(disk_name);
    if (!blk) {
        report_error("Disk not found: %s", disk_name);
        return;
    }

    BlockDriverState* bs = blk_bs(blk);
    if (!bs) {
        report_error("No block device attached to disk: %s", disk_name);
        return;
    }

    QCow2State *s = (QCow2State *)bs->opaque;
    if (!s) {
        report_error("Not a QCow2 file");
        return;
    }

    printf("\n=== QCow2 Analysis Report ===\n");
    printf("Disk Name: %s\n", disk_name);
    printf("File: %s\n", bs->filename);
    printf("Version: %d\n", s->version);
    printf("Virtual Size: %ld bytes (%ld MB)\n",
           bs->total_sectors * 512, (bs->total_sectors * 512) / (1024 * 1024));
    printf("Cluster Size: %d bytes\n", s->cluster_size);
    printf("Cluster Bits: %d\n", s->cluster_bits);
    printf("L2 Table Size: %d entries\n", s->l2_size);
    printf("L1 Table Size: %d entries\n", s->l1_size);
    printf("Crypt Method: %d\n", s->crypt_method);
    printf("Refcount Order: %d\n", s->refcount_order);

    if (s->version == QCOW_VERSION3) {
        printf("Incompatible Features: 0x%lx\n", s->incompatible_features);
        printf("Compatible Features: 0x%lx\n", s->compatible_features);
        printf("Autoclear Features: 0x%lx\n", s->autoclear_features);
    }

    char *backing = qcow2_get_backing_file(bs);
    if (backing) {
        printf("Backing File: %s\n", backing);
        qemu_free(backing);
    }
}

// 读取文件
static void read_file(const char *disk_name, int64_t offset, int64_t size, const char *output) {
    BlockBackend* blk = get_disk(disk_name);
    if (!blk) {
        report_error("Disk not found: %s", disk_name);
        return;
    }

    BlockDriverState* bs = blk_bs(blk);
    if (!bs) {
        report_error("No block device attached to disk: %s", disk_name);
        return;
    }

    uint8_t *buffer = (uint8_t *)qemu_malloc(size);
    if (bdrv_read(bs, offset, buffer, size) == size) {
        FILE *f = fopen(output, "wb");
        if (f) {
            fwrite(buffer, 1, size, f);
            fclose(f);
            report_info("Data written to %s", output);
        } else {
            report_error("Failed to write output file");
        }

        // 显示前64字节
        printf("First 64 bytes (hex):\n");
        hex_dump(buffer, (size < 64) ? size : 64, offset);

    } else {
        report_error("Failed to read data");
    }

    qemu_free(buffer);
}

// 写入文件
static void write_file(const char *disk_name, int64_t offset, const char *input_file,
                      const char *pattern, int64_t size, bool verify) {
    BlockBackend* blk = get_disk(disk_name);
    if (!blk) {
        report_error("Disk not found: %s", disk_name);
        return;
    }

    BlockDriverState* bs = blk_bs(blk);
    if (!bs) {
        report_error("No block device attached to disk: %s", disk_name);
        return;
    }

    std::vector<uint8_t> buffer;
    std::vector<uint8_t> verify_buffer;

    if (input_file) {
        // 从文件读取数据
        FILE *f = fopen(input_file, "rb");
        if (!f) {
            report_error("Failed to open input file: %s", input_file);
            return;
        }

        fseek(f, 0, SEEK_END);
        size_t file_size = ftell(f);
        fseek(f, 0, SEEK_SET);

        buffer.resize(file_size);
        if (fread(buffer.data(), 1, file_size, f) != file_size) {
            report_error("Failed to read input file");
            fclose(f);
            return;
        }
        fclose(f);

        if (verify) {
            verify_buffer.resize(file_size);
        }

        report_info("Writing %zu bytes from file %s at offset %ld",
                   file_size, input_file, offset);
    } else if (pattern) {
        // 使用模式填充
        if (size <= 0) {
            report_error("Size must be specified for pattern write");
            return;
        }

        buffer.resize(size);
        if (strcmp(pattern, "zero") == 0) {
            memset(buffer.data(), 0, size);
            report_info("Writing %ld zero bytes at offset %ld", size, offset);
        } else if (strncmp(pattern, "byte:", 5) == 0) {
            uint8_t byte_value = (uint8_t)strtoul(pattern + 5, NULL, 16);
            memset(buffer.data(), byte_value, size);
            report_info("Writing %ld bytes of 0x%02x at offset %ld",
                       size, byte_value, offset);
        } else if (strncmp(pattern, "pattern:", 8) == 0) {
            // 简单的模式填充
            const char *pat_str = pattern + 8;
            size_t pat_len = strlen(pat_str);
            for (int64_t i = 0; i < size; i++) {
                buffer[i] = pat_str[i % pat_len];
            }
            report_info("Writing %ld bytes with pattern '%s' at offset %ld",
                       size, pat_str, offset);
        } else {
            report_error("Unknown pattern: %s", pattern);
            return;
        }

        if (verify) {
            verify_buffer.resize(size);
        }
    } else {
        report_error("Either input file or pattern must be specified");
        return;
    }

    // 执行写操作
    ssize_t result = bdrv_write(bs, offset, buffer.data(), buffer.size());
    if (result == (ssize_t)buffer.size()) {
        report_info("Write completed successfully");

        // 验证写入
        if (verify) {
            report_info("Verifying write...");
            if (bdrv_read(bs, offset, verify_buffer.data(), verify_buffer.size()) == (ssize_t)verify_buffer.size()) {
                if (memcmp(buffer.data(), verify_buffer.data(), buffer.size()) == 0) {
                    report_info("Write verification passed");
                } else {
                    report_error("Write verification failed: data mismatch");
                }
            } else {
                report_error("Write verification failed: read error");
            }
        }
    } else {
        report_error("Write failed with error code: %d", result);
    }
}

// 查询数据块分布
static void query_data_blocks(const char *disk_name, bool fast_mode) {
    BlockBackend* blk = get_disk(disk_name);
    if (!blk) {
        report_error("Disk not found: %s", disk_name);
        return;
    }

    BlockDriverState* bs = blk_bs(blk);
    if (!bs) {
        report_error("No block device attached to disk: %s", disk_name);
        return;
    }

    auto start_time = get_timestamp_ms();

    std::list<DataBlockInfo*> block_list;
    if (fast_mode) {
        block_list = qcow2_query_data_blocks_fast(bs);
    } else {
        block_list = qcow2_query_data_blocks(bs);
    }

    auto end_time = get_timestamp_ms();
    auto duration = end_time - start_time;

    if (block_list.empty()) {
        report_error("Failed to query data blocks or no blocks found");
        return;
    }

    printf("\n=== Data Blocks Distribution ===\n");
    printf("Query time: %ld ms\n", duration);
    printf("%-12s %-12s %-12s %-12s %s\n",
           "Disk Offset", "Length", "File Offset", "Type", "Status");
    printf("------------ ------------ ------------ ------------ ---------\n");

    int total_blocks = 0;
    uint64_t total_allocated = 0;
    uint64_t total_unallocated = 0;

    for (DataBlockInfo* block : block_list) {
        total_blocks++;

        const char *type_str = "UNKNOWN";
        const char *status_str = "UNALLOCATED";

        switch (block->cluster_type) {
            case QCOW2_CLUSTER_NORMAL:
                type_str = "NORMAL";
                status_str = "ALLOCATED";
                total_allocated += block->length;
                break;
            case QCOW2_CLUSTER_COMPRESSED:
                type_str = "COMPRESSED";
                status_str = "ALLOCATED";
                total_allocated += block->length;
                break;
            case QCOW2_CLUSTER_UNALLOCATED:
                type_str = "UNALLOCATED";
                status_str = "UNALLOCATED";
                total_unallocated += block->length;
                break;
            case QCOW2_CLUSTER_ZERO:
                type_str = "ZERO";
                status_str = "ALLOCATED";
                total_allocated += block->length;
                break;
        }

        printf("0x%-10lx 0x%-10lx 0x%-10lx %-12s %s\n",
               (unsigned long)block->offset,
               (unsigned long)block->length,
               (unsigned long)block->file_offset,
               type_str,
               status_str);
    }

    printf("\n=== Summary ===\n");
    printf("Total blocks: %d\n", total_blocks);
    printf("Allocated space: %lu bytes (%lu MB)\n",
           total_allocated, total_allocated / (1024 * 1024));
    printf("Unallocated space: %lu bytes (%lu MB)\n",
           total_unallocated, total_unallocated / (1024 * 1024));

    uint64_t total_space = total_allocated + total_unallocated;
    if (total_space > 0) {
        printf("Allocation ratio: %.2f%%\n",
               (double)total_allocated / total_space * 100);
    } else {
        printf("Allocation ratio: 0.00%%\n");
    }

    // 清理
    for (DataBlockInfo* block : block_list) {
        qemu_free(block);
    }
}

// 创建文件
static void create_file(const char *filename, int64_t size, int cluster_size,
                       int version, const char *backing_file) {
    if (qcow2_create(filename, size, backing_file, cluster_size, version) == 0) {
        report_info("File created successfully");

        // 显示创建的文件信息
        BlockDriverState* bs = bdrv_open(filename, "qcow2");
        if (bs) {
            QCow2State *s = (QCow2State *)bs->opaque;
            if (s) {
                printf("\n=== Created QCow2 File Details ===\n");
                printf("Filename: %s\n", filename);
                printf("Version: %d\n", s->version);
                printf("Virtual Size: %ld bytes (%ld MB)\n",
                       bs->total_sectors * 512, (bs->total_sectors * 512) / (1024 * 1024));
                printf("Cluster Size: %d bytes\n", s->cluster_size);

                char *backing = qcow2_get_backing_file(bs);
                if (backing) {
                    printf("Backing File: %s\n", backing);
                    qemu_free(backing);
                } else {
                    printf("Backing File: None\n");
                }

                printf("Format: %s\n", (s->version == QCOW_VERSION3) ? "QCOW2 v3" : "QCOW2 v2");
            }
            bdrv_close(bs);
        }
    } else {
        report_error("Failed to create file");
    }
}

// 显示用法
static void print_usage(const char *program_name) {
    printf("QCow2 Tool - Production Ready QCow2 Disk Image Utility\n\n");
    printf("Usage:\n");
    printf("  %s open <name> <filename>          - Open QCow2 file with name\n", program_name);
    printf("  %s close <name>                    - Close named disk\n", program_name);
    printf("  %s list                            - List all open disks\n", program_name);
    printf("  %s analyze <name>                  - Analyze named disk\n", program_name);
    printf("  %s create <filename> -s <size> [-c <cluster_size>] [-v <version>] [-b <backing_file>]\n", program_name);
    printf("  %s read <name> -o <offset> -n <size> [-f <output>] - Read from named disk\n", program_name);
    printf("  %s write <name> -o <offset> [-i <input_file>] [-p <pattern>] [-n <size>] [-verify]\n", program_name);
    printf("  %s query <name> [-fast]            - Query data blocks distribution\n\n", program_name);

    printf("Create options:\n");
    printf("  -s <size>        Size in bytes (required)\n");
    printf("  -c <size>        Cluster size (default: 65536)\n");
    printf("  -v <version>     QCow2 version: 2 or 3 (default: 3)\n");
    printf("  -b <file>        Backing file path\n\n");

    printf("Write options:\n");
    printf("  -i <file>       Input file to write\n");
    printf("  -p <pattern>    Pattern: 'zero', 'byte:XX', or 'pattern:STR'\n");
    printf("  -n <size>       Size for pattern write\n");
    printf("  -verify         Verify after writing\n\n");

    printf("General options:\n");
    printf("  -o <offset>      Offset\n");
    printf("  -f <file>        Output file\n");
    printf("  -fast            Use fast mode for query\n\n");

    printf("Examples:\n");
    printf("  %s open base base.qcow2\n", program_name);
    printf("  %s open snap snapshot.qcow2\n", program_name);
    printf("  %s analyze base\n", program_name);
    printf("  %s read base -o 0 -n 512 -f boot.bin\n", program_name);
    printf("  %s create disk.qcow2 -s 1073741824 -v 3 -c 65536\n", program_name);
    printf("  %s create snap.qcow2 -s 10G -b base.qcow2\n", program_name);
    printf("  %s write disk -o 0 -i data.bin -verify\n", program_name);
    printf("  %s query disk -fast\n", program_name);
    printf("  %s list\n", program_name);
}

// 主函数
// 主函数
int main(int argc, char *argv[]) {
    // 如果提供了命令行参数，执行单次命令
    if (argc > 1) {
        const char *action = argv[1];
        bdrv_init();

        if (strcmp(action, "open") == 0) {
            if (argc < 4) {
                report_error("Usage: %s open <name> <filename>", argv[0]);
                return 1;
            }
            if (!open_disk(argv[2], argv[3])) {
                return 1;
            }
        } else if (strcmp(action, "close") == 0) {
            if (argc < 3) {
                report_error("Usage: %s close <name>", argv[0]);
                return 1;
            }
            close_disk(argv[2]);
        } else if (strcmp(action, "list") == 0) {
            list_disks();
        } else if (strcmp(action, "analyze") == 0) {
            if (argc < 3) {
                report_error("Usage: %s analyze <name>", argv[0]);
                return 1;
            }
            analyze_file(argv[2]);
        } else if (strcmp(action, "create") == 0) {
            const char *filename = NULL;
            int64_t size = 0;
            int cluster_size = 0;
            int version = 3;
            const char *backing_file = NULL;

            for (int i = 2; i < argc; i++) {
                if (strcmp(argv[i], "-s") == 0 && i + 1 < argc) {
                    size = atoll(argv[++i]);
                } else if (strcmp(argv[i], "-c") == 0 && i + 1 < argc) {
                    cluster_size = atoi(argv[++i]);
                } else if (strcmp(argv[i], "-v") == 0 && i + 1 < argc) {
                    version = atoi(argv[++i]);
                } else if (strcmp(argv[i], "-b") == 0 && i + 1 < argc) {
                    backing_file = argv[++i];
                } else if (filename == NULL) {
                    filename = argv[i];
                }
            }

            if (!filename || size <= 0) {
                report_error("Usage: %s create <filename> -s <size> [options]", argv[0]);
                return 1;
            }

            create_file(filename, size, cluster_size, version, backing_file);
        } else if (strcmp(action, "read") == 0) {
            const char *disk_name = NULL;
            int64_t offset = 0;
            int64_t size = 512;
            const char *output_file = "output.bin";

            for (int i = 2; i < argc; i++) {
                if (strcmp(argv[i], "-o") == 0 && i + 1 < argc) {
                    offset = atoll(argv[++i]);
                } else if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) {
                    size = atoll(argv[++i]);
                } else if (strcmp(argv[i], "-f") == 0 && i + 1 < argc) {
                    output_file = argv[++i];
                } else if (disk_name == NULL) {
                    disk_name = argv[i];
                }
            }

            if (!disk_name) {
                report_error("Usage: %s read <name> -o <offset> -n <size> [-f <output>]", argv[0]);
                return 1;
            }

            read_file(disk_name, offset, size, output_file);
        } else if (strcmp(action, "write") == 0) {
            const char *disk_name = NULL;
            int64_t offset = 0;
            const char *input_file = NULL;
            const char *pattern = NULL;
            int64_t size = 0;
            bool verify = false;

            for (int i = 2; i < argc; i++) {
                if (strcmp(argv[i], "-o") == 0 && i + 1 < argc) {
                    offset = atoll(argv[++i]);
                } else if (strcmp(argv[i], "-i") == 0 && i + 1 < argc) {
                    input_file = argv[++i];
                } else if (strcmp(argv[i], "-p") == 0 && i + 1 < argc) {
                    pattern = argv[++i];
                } else if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) {
                    size = atoll(argv[++i]);
                } else if (strcmp(argv[i], "-verify") == 0) {
                    verify = true;
                } else if (disk_name == NULL) {
                    disk_name = argv[i];
                }
            }

            if (!disk_name || offset < 0) {
                report_error("Usage: %s write <name> -o <offset> [-i <input_file>] [-p <pattern>] [-n <size>] [-verify]", argv[0]);
                return 1;
            }

            if (!input_file && !pattern) {
                report_error("Either input file (-i) or pattern (-p) must be specified");
                return 1;
            }

            if (pattern && size <= 0) {
                report_error("Size must be specified with -n for pattern write");
                return 1;
            }

            write_file(disk_name, offset, input_file, pattern, size, verify);
        } else if (strcmp(action, "query") == 0) {
            const char *disk_name = NULL;
            bool fast_mode = false;

            for (int i = 2; i < argc; i++) {
                if (strcmp(argv[i], "-fast") == 0) {
                    fast_mode = true;
                } else if (disk_name == NULL) {
                    disk_name = argv[i];
                }
            }

            if (!disk_name) {
                report_error("Usage: %s query <name> [-fast]", argv[0]);
                return 1;
            }

            query_data_blocks(disk_name, fast_mode);
        } else if (strcmp(action, "help") == 0) {
            print_usage(argv[0]);
        } else {
            report_error("Unknown action: %s", action);
            print_usage(argv[0]);
            return 1;
        }

        return 0;
    }

    // 交互式模式
    printf("QCow2 Tool - Interactive Mode\n");
    printf("Type 'help' for usage, 'exit' or 'quit' to exit\n\n");

    bdrv_init();

    char line[1024];
    while (1) {
        printf("qcow2> ");
        fflush(stdout);

        if (!fgets(line, sizeof(line), stdin)) {
            break;
        }

        // 去除换行符
        line[strcspn(line, "\n")] = 0;

        // 跳过空行
        if (strlen(line) == 0) {
            continue;
        }

        // 退出命令
        if (strcmp(line, "exit") == 0 || strcmp(line, "quit") == 0) {
            break;
        }

        // 帮助命令
        if (strcmp(line, "help") == 0) {
            print_usage("qcow2_tool");
            continue;
        }

        // 解析命令行
        int interactive_argc = 0;
        char *interactive_argv[64];
        char *token = strtok(line, " ");

        while (token != NULL && interactive_argc < 63) {
            interactive_argv[interactive_argc++] = token;
            token = strtok(NULL, " ");
        }
        interactive_argv[interactive_argc] = NULL;

        if (interactive_argc == 0) {
            continue;
        }

        // 执行命令
        const char *action = interactive_argv[0];

        if (strcmp(action, "open") == 0) {
            if (interactive_argc < 3) {
                report_error("Usage: open <name> <filename>");
                continue;
            }
            if (!open_disk(interactive_argv[1], interactive_argv[2])) {
                continue;
            }
        } else if (strcmp(action, "close") == 0) {
            if (interactive_argc < 2) {
                report_error("Usage: close <name>");
                continue;
            }
            close_disk(interactive_argv[1]);
        } else if (strcmp(action, "list") == 0) {
            list_disks();
        } else if (strcmp(action, "analyze") == 0) {
            if (interactive_argc < 2) {
                report_error("Usage: analyze <name>");
                continue;
            }
            analyze_file(interactive_argv[1]);
        } else if (strcmp(action, "create") == 0) {
            const char *filename = NULL;
            int64_t size = 0;
            int cluster_size = 0;
            int version = 3;
            const char *backing_file = NULL;

            for (int i = 1; i < interactive_argc; i++) {
                if (strcmp(interactive_argv[i], "-s") == 0 && i + 1 < interactive_argc) {
                    size = atoll(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-c") == 0 && i + 1 < interactive_argc) {
                    cluster_size = atoi(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-v") == 0 && i + 1 < interactive_argc) {
                    version = atoi(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-b") == 0 && i + 1 < interactive_argc) {
                    backing_file = interactive_argv[++i];
                } else if (filename == NULL) {
                    filename = interactive_argv[i];
                }
            }

            if (!filename || size <= 0) {
                report_error("Usage: create <filename> -s <size> [options]");
                continue;
            }

            create_file(filename, size, cluster_size, version, backing_file);
        } else if (strcmp(action, "read") == 0) {
            const char *disk_name = NULL;
            int64_t offset = 0;
            int64_t size = 512;
            const char *output_file = "output.bin";

            for (int i = 1; i < interactive_argc; i++) {
                if (strcmp(interactive_argv[i], "-o") == 0 && i + 1 < interactive_argc) {
                    offset = atoll(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-n") == 0 && i + 1 < interactive_argc) {
                    size = atoll(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-f") == 0 && i + 1 < interactive_argc) {
                    output_file = interactive_argv[++i];
                } else if (disk_name == NULL) {
                    disk_name = interactive_argv[i];
                }
            }

            if (!disk_name) {
                report_error("Usage: read <name> -o <offset> -n <size> [-f <output>]");
                continue;
            }

            read_file(disk_name, offset, size, output_file);
        } else if (strcmp(action, "write") == 0) {
            const char *disk_name = NULL;
            int64_t offset = 0;
            const char *input_file = NULL;
            const char *pattern = NULL;
            int64_t size = 0;
            bool verify = false;

            for (int i = 1; i < interactive_argc; i++) {
                if (strcmp(interactive_argv[i], "-o") == 0 && i + 1 < interactive_argc) {
                    offset = atoll(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-i") == 0 && i + 1 < interactive_argc) {
                    input_file = interactive_argv[++i];
                } else if (strcmp(interactive_argv[i], "-p") == 0 && i + 1 < interactive_argc) {
                    pattern = interactive_argv[++i];
                } else if (strcmp(interactive_argv[i], "-n") == 0 && i + 1 < interactive_argc) {
                    size = atoll(interactive_argv[++i]);
                } else if (strcmp(interactive_argv[i], "-verify") == 0) {
                    verify = true;
                } else if (disk_name == NULL) {
                    disk_name = interactive_argv[i];
                }
            }

            if (!disk_name || offset < 0) {
                report_error("Usage: write <name> -o <offset> [-i <input_file>] [-p <pattern>] [-n <size>] [-verify]");
                continue;
            }

            if (!input_file && !pattern) {
                report_error("Either input file (-i) or pattern (-p) must be specified");
                continue;
            }

            if (pattern && size <= 0) {
                report_error("Size must be specified with -n for pattern write");
                continue;
            }

            write_file(disk_name, offset, input_file, pattern, size, verify);
        } else if (strcmp(action, "query") == 0) {
            const char *disk_name = NULL;
            bool fast_mode = false;

            for (int i = 1; i < interactive_argc; i++) {
                if (strcmp(interactive_argv[i], "-fast") == 0) {
                    fast_mode = true;
                } else if (disk_name == NULL) {
                    disk_name = interactive_argv[i];
                }
            }

            if (!disk_name) {
                report_error("Usage: query <name> [-fast]");
                continue;
            }

            query_data_blocks(disk_name, fast_mode);
        } else {
            report_error("Unknown command: %s", action);
            printf("Available commands: open, close, list, analyze, create, read, write, query, help, exit\n");
        }
    }

    // 清理所有打开的磁盘
    printf("Cleaning up...\n");
    while (!blk_map.empty()) {
        auto it = blk_map.begin();
        printf("Closing disk: %s\n", it->first.c_str());
        blk_unref(it->second);
        blk_map.erase(it);
    }

    printf("Goodbye!\n");
    return 0;
}